import torch
import numpy as np
import cv2


def test_numpy_tensor():
    print("=========================== numpy <-> tensor =====")
    arr = np.arange(10).reshape(2, 5)
    print("np.arange(10).reshape(2, 5): ")
    print(arr)
    print()

    a = torch.from_numpy(arr)
    print("torch.from_numpy(arr): ")
    print(a)
    print()

    print("a.numpy(): ")
    out = a.numpy()
    print(out)
    print()


def test_cv2():
    print("================ cv2 to_cuda numpy <-> tensor ====")
    data = cv2.imread("test.png")
    print(data)
    print()
    cv2.imshow("test1", data)

    print(" numpy -->  tensor")
    out = torch.from_numpy(data)
    print(out)
    print()

    print(" cpu --> gpu")
    out = out.to(torch.device("cuda"))
    print("out.is_cuda: ", out.is_cuda)

    print(" flip 翻转")
    out = torch.flip(out, dims=[0])

    print(" gpu --> cpu")
    out = out.to(torch.device("cpu"))
    print("out.is_cuda: ", out.is_cuda)

    print(" tensor --> numpy")
    data = out.numpy()
    print(data)

    cv2.imshow("test2", data)
    cv2.waitKey(0)


if __name__ == '__main__':
    test_numpy_tensor()
    test_cv2()
